import numpy as np



# 初始化参数
num_ants = 10
num_resources = 5
num_demands = 5
pheromone = np.ones((num_resources, num_demands))
visibility = np.random.rand(num_resources, num_demands)
alpha = 1
beta = 2
evaporation = 0.5

# 更新信息素
def update_pheromone(pheromone, delta_pheromone):
    pheromone = (1 - evaporation) * pheromone + delta_pheromone
    return pheromone

# 蚁群算法
def ant_colony(num_ants, num_iterations):
    global pheromone
    best_solution = None
    best_cost = float('inf')

    for _ in range(num_iterations):
        solutions = []
        costs = []

        for _ in range(num_ants):
            solution = np.zeros((num_resources, num_demands))
            cost = 0

            # 核心内容
            for i in range(num_resources):
                probabilities = (pheromone[i] ** alpha) * (visibility[i] ** beta)
                probabilities /= np.sum(probabilities)
                demand = np.random.choice(num_demands, p=probabilities)
                solution[i][demand] = 1
                cost += visibility[i][demand]

            solutions.append(solution)
            costs.append(cost)

            # 更新最优解
            if cost < best_cost:
                best_solution = solution
                best_cost = cost

        # 更新信息素
        delta_pheromone = np.zeros((num_resources, num_demands))
        for solution, cost in zip(solutions, costs):
            delta_pheromone += solution * (1 / cost)

        pheromone = update_pheromone(pheromone, delta_pheromone)

    return best_solution, best_cost

if __name__ == '__main__':
    # 测试算法
    best_solution, best_cost = ant_colony(num_ants=10, num_iterations=100)
    print("Best solution:", best_solution)
    print("Best cost:", best_cost)
